package com.log.enhance.util;

import com.log.enhance.constant.LogPluginConstant;
import org.slf4j.MDC;

import java.util.Map;
import java.util.concurrent.Callable;

/**
 * 日志多线程情况公用工具类
 * 更优雅的实现参考：
 * https://www.jianshu.com/p/151d13011112
 * https://gitee.com/alibaba-projects/transmittable-thread-local
 *
 * @author Mr_wenpan@163.com 2021/05/29 11:45
 */
public class ThreadMdcUtil {

    /**
     * 设置traceId到MDC中（如果MDC中不存在对应的traceId的话）
     *
     * @return void
     */
    public static void setTraceIdIfAbsent() {
        if (MDC.get(LogPluginConstant.TRACEID) == null) {
            MDC.put(LogPluginConstant.TRACEID, LogCommonUtil.generateTraceId());
        }
    }

    /**
     * 设置traceId到MDC中（直接通过随机生成的traceId覆盖原有的traceId）
     *
     * @return void
     */
    public static void setTraceId() {
        MDC.put(LogPluginConstant.TRACEID, LogCommonUtil.generateTraceId());
    }

    /**
     * 设置traceId到MDC中（直接通过传入的traceId覆盖原有的traceId）
     *
     * @param traceId traceId
     * @return void
     */
    public static void setTraceId(String traceId) {
        MDC.put(LogPluginConstant.TRACEID, traceId);
    }

    /**
     * 包装callable
     *
     * @param callable 线程任务
     * @param context  创建任务的线程传递给执行任务的线程的数据map
     * @return java.util.concurrent.Callable<T>
     * @date 2021/6/1 3:09 下午
     * @author wenpan
     */
    public static <T> Callable<T> wrap(Callable<T> callable, Map<String, String> context) {
        return () -> {
            // 实际执行前导入对应请求的MDC副本
            if (context == null) {
                MDC.clear();
            } else {
                // 导入请求对应的Map副本到MDC中
                MDC.setContextMap(context);
            }
            // 如果该线程中没有traceId则生成一个traceId并且放入到MDC中
            setTraceIdIfAbsent();
            try {
                // 调用真正的callable任务执行
                return callable.call();
            } finally {
                // 清除MDC内容，防止线程池线程重用的干扰
                MDC.clear();
            }
        };
    }

    /**
     * 包装 runnable
     *
     * @param runnable 线程任务
     * @param context  创建任务的线程传递给执行任务的线程的数据map
     * @return java.util.concurrent.Callable<T>
     * @date 2021/6/1 3:19 下午
     * @author wenpan
     */
    public static Runnable wrap(Runnable runnable, Map<String, String> context) {
        return () -> {
            // 如果context为空，则说明没有数据需要传递，直接清空MDC中的threadlocal中的数据即可
            // 否则将数据传递到执行任务的线程中
            if (context == null) {
                MDC.clear();
            } else {
                MDC.setContextMap(context);
            }
            // 如果该线程中没有traceId则生成一个traceId并且放入到MDC中
            setTraceIdIfAbsent();
            try {
                // 真正执行runnable线程任务
                runnable.run();
            } finally {
                // 任务执行完毕后需要清空threadlocal，防止线程池线程重用造成干扰
                MDC.clear();
            }
        };
    }

}

